#ifndef DEEPDR_CLASSIFIER_H
#define DEEPDR_CLASSIFIER_H

#include <opencv/cv.h>
#include <opencv/highgui.h>
#include <opencv2/opencv.hpp>
#include <vector>
#include <string>
#include <map>
#include <lmdb.h>
#include <caffe/proto/caffe.pb.h>
#include <glog/logging.h>
#include <gtest/gtest.h>
#include "common.h"
#include "caffe/caffe.hpp"

using cv::Mat;
using std::vector;
using std::string;
using std::pair;
using std::map;

struct Sample {
    Mat data;
    int label;
    int index;
};

namespace Parameters {
    const int NUM_LABELS = 1;
    const int NUM_STRIDE = 2;
    const int IMAGE_COLS = 1200;
    const int WINDOW_SIZE = 41;
};

class CommonImage;

void MDB_CHECK(int mdb_status);

class CommonImage {
    Mat originalImage;
protected:
    virtual string registerFuncName(string) final;

    map<string, CommonImage> imgStorage;
    map<string, int> registeredNames;
    string name;
public:
    CommonImage(Mat img, string name = "");

    CommonImage(string name = "");

    void setImage(Mat img);

    virtual CommonImage &getDouble();

    virtual CommonImage &getByte();

    virtual CommonImage &getGreenChannel();

    virtual CommonImage &getChannel();

    virtual CommonImage &getLocalNorm(
            float sigma1,
            float sigma2,
            cv::InputArray mask = cv::noArray());

    virtual const Mat &getOriginal() const;

    virtual void show(string name = "", int pause = -1) const;

    virtual void write(int id = -1) const;
};


class DRImage : public CommonImage {
    int id;
    static int NUM_DRIMAGE;
public:
    DRImage(int _id = -1);

    DRImage(Mat img, int _id = -1);

    const CommonImage &getMask(double thresh = 0.001);

    const CommonImage &getHardLeakageNoneCaffe();

    const CommonImage &getVesselNoneCaffe();

    const CommonImage &getIgnoreFiel();

    const CommonImage &getGrowBoundray();
};


/*
 * The Virtual Class DataSet
 * Provide a size and subscript operator
 * ImageData form original Files and LMDB datset
 * should dirive from this class
 */
class DataSet {
public:
    virtual Sample operator[](int index) = 0;

    virtual size_t length() const = 0;
};

/*
 *  This class is designed to load data form LMDB
 *  database, And convert to Mat in OpenCV or get
 *  original Datum.
 *  Thus two things need to be done.
 *  First, there seems no Datum to Mat convert
 *  code in caffe, So I need to implement it.
 *  In order to check the correctness of my code,
 *  Tests needs to be added to Unit test.
 *
 *  interface should interact with eachother
 *  Maybe I should combine these two interface by
 *  calling cursor interface in index implementation;
 */

class Transaction {
protected:
    virtual void open() = 0;

    virtual void close() = 0;

public:
    virtual void commit() = 0;

    virtual string get(const string &key) = 0;

    virtual void put(const string &key, const string &val) = 0;
};

class ReadTransaction : Transaction {
private:
    MDB_env *env;
    MDB_txn *txn;
    MDB_dbi dbi;

public:
    ReadTransaction(MDB_env *mdb_env_) : env(mdb_env_) {
        open();
    }

    ~ReadTransaction() {
        close();
    }

protected:
    virtual void open() override {
        CHECK(env != NULL) << "Error Open DB" << Common::printTrace();
        MDB_CHECK(mdb_txn_begin(env, NULL, MDB_RDONLY, &txn));
        MDB_CHECK(mdb_dbi_open(txn, NULL, 0, &dbi));
    }

    virtual void close() override {
        if (txn != NULL)
            mdb_txn_abort(txn);
        if (dbi)
            mdb_dbi_close(env, dbi);
        txn = NULL;
        dbi = 0;
    }

public:
    virtual void commit() override {
        LOG(WARNING) <<
        "should not call this function for read only Transaction";
    }

    virtual void put(const string &key, const string &val) override {
        LOG(FATAL) << Common::printTrace() <<
        "Should not write to a read only transtraction";
    }

    virtual string get(const string &key) override {
        MDB_val mdb_key_, mdb_data;
        mdb_key_.mv_size = key.size();
        mdb_key_.mv_data = (void *) key.data();
        int res = mdb_get(txn, dbi, &mdb_key_, &mdb_data);
        if (res == MDB_NOTFOUND) {
            LOG(ERROR) << "Error reading key : " << key;
            return "";
        } else {
            MDB_CHECK(res);
        }
        return string((char *) mdb_data.mv_data, mdb_data.mv_size);
    }
};

class WriteTransaction : Transaction {
private:
    MDB_env *env;
    MDB_txn *txn;
    MDB_dbi dbi;
    vector<string> keys;
    vector<string> values;

protected:
    void open() {
        CHECK(env != NULL) << "Error Open DB" << Common::printTrace();
        MDB_CHECK(mdb_txn_begin(env, NULL, MDB_WRITEMAP, &txn));
        MDB_CHECK(mdb_dbi_open(txn, NULL, MDB_CREATE, &dbi));
    }

    void close() {
        if (txn != NULL)
            mdb_txn_abort(txn);
        if (dbi)
            mdb_dbi_close(env, dbi);
        txn = NULL;
        dbi = 0;
    }

public:
    void commit() {
        open();
        MDB_val mdb_key, mdb_data;
        bool out_of_memory = false;
        for (int i = 0; i < keys.size(); i++) {
            mdb_key.mv_size = keys[i].size();
            mdb_key.mv_data = const_cast<char *>(keys[i].data());
            mdb_data.mv_size = values[i].size();
            mdb_data.mv_data = const_cast<char *>(values[i].data());
            int put_rc = mdb_put(txn, dbi, &mdb_key, &mdb_data, 0);
            if (put_rc == MDB_MAP_FULL) {
                out_of_memory = true;
                break;
            } else {
                // Failed for some other reason
                MDB_CHECK(put_rc);
            }
        }
        if (!out_of_memory) {
            // Commit the transaction
            MDB_CHECK(mdb_txn_commit(txn));
            mdb_dbi_close(env, dbi);
            keys.clear();
            values.clear();
        } else {
            // Double the map size and retry
            close();
            doubleSpace();
            commit();
        }
    }

    WriteTransaction(MDB_env *mdb_env_) :
            env(mdb_env_) {
    }

    ~WriteTransaction() {
        close();
    }

    void put(const string &key, const string &val) {
        keys.push_back(key);
        values.push_back(val);
    }

    string get(const string &key) {
        CHECK(false) << Common::printTrace() <<
                     "Should Not Reach this function";
        return "";
    }

    void doubleSpace() {
        struct MDB_envinfo current_info;
        MDB_CHECK(mdb_env_info(env, &current_info));
        size_t new_size = current_info.me_mapsize * 2;
        MDB_CHECK(mdb_env_set_mapsize(env, new_size));
        LOG(INFO) << "Double space to " << (double) new_size / 1024.0 / 1024.0;
    }
};

class LMDBDataSet {

private:
    MDB_env *mdb_env_;
    size_t size;
    size_t keySize;
    map<int, bool> existedKeys;

    Transaction *transaction;

    int writeCounter;
private:

    static Transaction *newWriteTransaction(MDB_env *env) {
        return (Transaction *) new WriteTransaction(env);
    }

    static Transaction *newReadTransaction(MDB_env *env) {
        return (Transaction *) new ReadTransaction(env);
    }


private:
    int currentPosition;
    bool valid_;

    inline size_t getKeySize();

    inline void put(string &key, string &val) {
        const static int commitIteration = 100000;
        transaction->put(key, val);
        writeCounter++;
        if (writeCounter % commitIteration == commitIteration - 1) {
            transaction->commit();
        }
    }

public:
    LMDBDataSet();

    virtual  ~LMDBDataSet();

    virtual bool loadFromLMDB(string DBPath);

    virtual void newDB(string DBPath);

    virtual void addSample(Sample sample, int index = -1);

    virtual void addDatum(int key, caffe::Datum *datum);

    virtual void close();

    virtual Sample operator[](int index);

    virtual size_t length();

    /*
     * Below is the cursor like interface;
     *
     * Implement a cursor like interface.
     * I have to do this to aviod extra data copy
     * and extra database operation.
     */
    virtual bool seek(int index);

    virtual string getMDB();

    virtual bool getDatum(caffe::Datum &datum);

    virtual Mat getCvMat();

    virtual Sample getSample();

    static Mat datumToMat(caffe::Datum *datum);
};

class ImgDataGenerator : DataSet {
private:
    class ImageSet {
    public:
        DRImage originalImage;
        vector<CommonImage> labels;
    };

    struct DataSample {
        int imgIndex;
        int row, col;
        int labels[Parameters::NUM_LABELS];
    };

private:
    vector<ImageSet> images;
    vector<DataSample> dataSamples;

public:
    ImgDataGenerator(string codeFilePath);

    virtual Sample operator[](int index) override;

    virtual size_t length() const override;
};


class Classifier {
public:
    ImgDataGenerator dg; // including all data avail;
protected:
    // reset sample and meke new samples;
//    virtual void resetSample();
//    virtual void train();
public:
    Classifier(string trainFileList);
//    void learn();
//    virtual vector<Mat> predict(vector<Mat> data);
};


#endif //DEEPDR_CLASSIFIER_H
